!> author: 左志华
!> date: 2022-05-31
!>
!> 通用数学函数库
module forlab_m

    use, intrinsic :: ieee_arithmetic, only: ieee_is_nan
    use, intrinsic :: iso_c_binding, only: c_null_char
    implicit none
    private

    public :: arg, argd, cross_product, diff, is_close, Pi
    public :: mean, var, cov, randu, randn
    public :: to_string, to_lower, to_upper, f_c_string, progress_bar, puts
    public :: newline
    public :: starts_with
    public :: error_stop
    public :: is_exist, mkdir, rmdir, remove, &
              is_windows, pwd, rename, countlines, windows_path
    public :: tic, toc, time_stamp, is_leap
    public :: arange, diag, eye, linspace, tri, tril, triu

    interface diff
        module procedure :: diff_real, diff_integer
    end interface diff

    interface is_close
        module procedure :: is_close_real, is_close_cmplx
    end interface is_close

    real, parameter :: Pi = acos(-1.0)
    real, parameter :: sqrt_eps = sqrt(epsilon(1.0))

    !> 平均数
    interface mean
        module procedure :: mean_real, mean_integer
    end interface mean

    !> 方差
    interface var
        module procedure :: var_real, var_integer
    end interface var

    !> 生成 [0,1) 随机数
    interface randu
        module procedure :: randu_rank1, randu_rank0
    end interface randu

    !> 生成 N(0,1) 正态分布随机数
    interface randn
        module procedure :: randn_rank1, randn_rank0
    end interface randn

    !> 将其他类型转化为字符串
    interface to_string
        module procedure to_string_integer, to_string_logical, to_string_real
    end interface to_string

    character(*), parameter :: uppercase = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    character(*), parameter :: lowercase = "abcdefghijklmnopqrstuvwxyz"
    character(*), parameter :: newline = new_line('')

    interface error_stop
        module procedure :: error_stop_msg, error_stop_code
    end interface error_stop

    integer, save :: time_save !! save the time

    interface arange
        module procedure :: arange_integer, arange_real
    end interface arange

    interface diag
        module procedure :: diag_real_rank2, diag_real_rank1
    end interface diag

    interface linspace
        module procedure :: linspace_integer, linspace_real
    end interface linspace

contains

    !> calculate single precision cross product
    pure function cross_product(a, b) result(c)
        real, intent(in) :: a(3), b(3)
        real :: c(3)

        c(1) = a(2)*b(3) - a(3)*b(2)
        c(2) = a(3)*b(1) - a(1)*b(3)
        c(3) = a(1)*b(2) - a(2)*b(1)

    end function cross_product

    !> calculate argument of complex number
    elemental function arg(z) result(a)
        complex, intent(in) :: z
        real :: a

        if (z == (0.0, 0.0)) then
            a = 0.0
        else
            a = atan2(z%im, z%re)
        end if
    end function arg

    !> calculate argument of complex number
    elemental function argd(z) result(a)
        complex, intent(in) :: z
        real :: a

        if (z == (0.0, 0.0)) then
            a = 0.0
        else
            a = atan2(z%im, z%re)*180.0/Pi
        end if
    end function argd

    !> calculate difference of x with n elements
    pure function diff_real(x, n, prepend, append) result(y)
        real, intent(in) :: x(:)
        integer, intent(in), optional :: n
        real, intent(in), optional :: prepend(:), append(:)
        real, allocatable :: y(:)
        integer :: size_prepend, size_append, size_x, size_work
        integer :: n_, i

        if (present(n)) then
            n_ = n
        else
            n_ = 1
        end if

        if (n_ <= 0) then
            y = x
            return
        end if

        size_prepend = 0
        size_append = 0
        if (present(prepend)) size_prepend = size(prepend)
        if (present(append)) size_append = size(append)
        size_x = size(x)
        size_work = size_x + size_prepend + size_append

        if (size_work <= n_) then
            allocate (y(0))
            return
        end if

        !> Use a quick exit for the common case, to avoid memory allocation.
        if (size_prepend == 0 .and. size_append == 0 .and. n_ == 1) then
            y = x(2:) - x(1:size_x - 1)
            return
        end if

        block
            real :: work(size_work)
            if (size_prepend > 0) work(:size_prepend) = prepend
            work(size_prepend + 1:size_prepend + size_x) = x
            if (size_append > 0) work(size_prepend + size_x + 1:) = append

            do i = 1, n_
                work(1:size_work - i) = work(2:size_work - i + 1) - work(1:size_work - i)
            end do

            y = work(1:size_work - n_)
        end block
    end function diff_real

    !> calculate difference of x with n elements
    pure function diff_integer(x, n, prepend, append) result(y)
        integer, intent(in) :: x(:)
        integer, intent(in), optional :: n
        integer, intent(in), optional :: prepend(:), append(:)
        integer, allocatable :: y(:)
        integer :: size_prepend, size_append, size_x, size_work
        integer :: n_, i

        if (present(n)) then
            n_ = n
        else
            n_ = 1
        end if

        if (n_ <= 0) then
            y = x
            return
        end if

        size_prepend = 0
        size_append = 0
        if (present(prepend)) size_prepend = size(prepend)
        if (present(append)) size_append = size(append)
        size_x = size(x)
        size_work = size_x + size_prepend + size_append

        if (size_work <= n_) then
            allocate (y(0))
            return
        end if

        !> Use a quick exit for the common case, to avoid memory allocation.
        if (size_prepend == 0 .and. size_append == 0 .and. n_ == 1) then
            y = x(2:) - x(1:size_x - 1)
            return
        end if

        block
            integer :: work(size_work)
            if (size_prepend > 0) work(:size_prepend) = prepend
            work(size_prepend + 1:size_prepend + size_x) = x
            if (size_append > 0) work(size_prepend + size_x + 1:) = append

            do i = 1, n_
                work(1:size_work - i) = work(2:size_work - i + 1) - work(1:size_work - i)
            end do

            y = work(1:size_work - n_)
        end block
    end function diff_integer

    !> check if a and b are close
    elemental logical function is_close_real(a, b, rel_tol, abs_tol, equal_nan) result(close)
        real, intent(in) :: a, b
        real, intent(in), optional :: rel_tol, abs_tol
        logical, intent(in), optional :: equal_nan
        logical :: equal_nan_
        real :: abs_tol_, rel_tol_

        if (present(equal_nan)) then
            equal_nan_ = equal_nan
        else
            equal_nan_ = .false.
        end if

        if (ieee_is_nan(a) .or. ieee_is_nan(b)) then
            close = merge(.true., .false., equal_nan_ .and. ieee_is_nan(a) .and. ieee_is_nan(b))
        else

            if (present(rel_tol)) then
                rel_tol_ = rel_tol
            else
                rel_tol_ = sqrt_eps
            end if

            if (present(abs_tol)) then
                abs_tol_ = abs_tol
            else
                abs_tol_ = 0.0
            end if

            close = abs(a - b) <= max(abs(rel_tol_*max(abs(a), abs(b))), &
                                      abs(abs_tol_))
        end if

    end function is_close_real

    !> check if a and b are close
    elemental logical function is_close_cmplx(a, b, rel_tol, abs_tol, equal_nan) result(close)
        complex, intent(in) :: a, b
        real, intent(in), optional :: rel_tol, abs_tol
        logical, intent(in), optional :: equal_nan

        close = is_close_real(a%re, b%re, rel_tol, abs_tol, equal_nan) .and. &
                is_close_real(a%im, b%im, rel_tol, abs_tol, equal_nan)
    end function is_close_cmplx

    !> Compute the mean of a vector of single precision real numbers
    pure function mean_real(x) result(y)
        real, intent(in) :: x(:)
        real :: y
        y = sum(x)/size(x)
    end function mean_real

    !> Compute the mean of a vector of integer numbers
    pure function mean_integer(x) result(y)
        integer, intent(in) :: x(:)
        integer :: y
        y = sum(x)/size(x)
    end function mean_integer

    !> Compute the variance of a vector of single precision real numbers
    pure function var_real(x, normalized) result(y)
        real, intent(in) :: x(:)
        logical, optional, intent(in) :: normalized !! by default, .false.
        real :: y
        logical :: normalized_

        if (present(normalized)) then
            normalized_ = normalized
        else
            normalized_ = .false.
        end if

        if (normalized_) then
            y = sum((x - mean(x))**2)/(size(x) - 1)
        else
            y = sum((x - mean(x))**2)/size(x)
        end if

    end function var_real

    !> Compute the variance of a vector of integer numbers
    pure function var_integer(x, normalized) result(y)
        integer, intent(in) :: x(:)
        logical, optional, intent(in) :: normalized !! by default, .false.
        integer y
        logical normalized_

        if (present(normalized)) then
            normalized_ = normalized
        else
            normalized_ = .false.
        end if

        if (normalized_) then
            y = sum((x - mean_integer(x))**2)/(size(x) - 1)
        else
            y = sum((x - mean_integer(x))**2)/size(x)
        end if

    end function var_integer

    !> generates a vector of randon numbers uniformly distributed in [0, 1).
    impure function randu_rank1(n) result(a)
        integer, intent(in) :: n
        real :: a(n)

        call random_number(a)
    end function randu_rank1

    !> generates a random number uniformly distributed in [0, 1).
    impure function randu_rank0() result(x)
        real :: x

        call random_number(x)
    end function randu_rank0

    !> generates a random number normally distributed with mean 0 and variance 1.
    impure function randn_rank0() result(a)
        real a
        real u, v

        do
            call random_number(u)
            call random_number(v)
            u = 2*u - 1
            v = 2*v - 1
            a = u*u + v*v
            if (a > 0 .and. a < 1) exit
        end do
        a = u*sqrt(-2*log(a)/a)
    end function randn_rank0

    !> generates a vector of randon numbers normally distributed with mean 0 and standard deviation 1.
    impure function randn_rank1(n) result(a)
        integer, intent(in) :: n
        real a(n)
        integer i

        do i = 1, n
            a(i) = randn_rank0()
        end do
    end function randn_rank1

    !> Compute the covariance of a vector of single precision real numbers
    pure function cov(x, normalized) result(y)
        real, intent(in) :: x(:)
        logical, optional, intent(in) :: normalized !! by default, .false.
        real y
        logical normalized_
        real x_(size(x))

        if (present(normalized)) then
            normalized_ = normalized
        else
            normalized_ = .false.
        end if

        x_ = x - mean(x)
        if (normalized_) then
            y = dot_product(x_, x_)/(size(x) - 1)
        else
            y = dot_product(x_, x_)/size(x)
        end if

    end function cov

    !> Make a string lowercase
    elemental function to_lower(x) result(y)
        character(*), intent(in) :: x   !! Input string
        character(len(x)) y
        integer i, k

        do concurrent(i=1:len(x))
            k = index(uppercase, x(i:i))
            if (k > 0) then
                y(i:i) = lowercase(k:k)
            else
                y(i:i) = x(i:i)
            end if
        end do

    end function to_lower

    !> Make a string uppercase
    elemental function to_upper(x) result(y)
        character(*), intent(in) :: x   !! Input string
        character(len(x)) y
        integer i, k

        do concurrent(i=1:len(x))
            k = index(lowercase, x(i:i))
            if (k > 0) then
                y(i:i) = uppercase(k:k)
            else
                y(i:i) = x(i:i)
            end if
        end do

    end function to_upper

    !> Format an integer to a string
    pure function to_string_integer(i, fmt) result(s)
        integer, intent(in) :: i                    !! integer to format
        character(*), intent(in), optional :: fmt   !! format string
        character(:), allocatable :: s              !! resulting string
        character(20) :: s_

        if (present(fmt)) then
            write (s_, "("//fmt//")") i
        else
            write (s_, *) i
        end if
        s = trim(s_)
    end function to_string_integer

    !> Format a logical to a string
    pure function to_string_logical(l, fmt) result(s)
        logical, intent(in) :: l                    !! logical to format
        character(*), intent(in), optional :: fmt   !! format string
        character(10) :: s_
        character(:), allocatable :: s

        if (present(fmt)) then
            write (s_, "("//fmt//")") l
            s = trim(s_)
        else
            if (l) then
                s = "T"
            else
                s = "F"
            end if
        end if
    end function to_string_logical

    !> Format a real number to a string
    pure function to_string_real(r, fmt) result(s)
        real, intent(in) :: r                       !! real to format
        character(*), intent(in), optional :: fmt   !! format string
        character(:), allocatable :: s
        character(128) :: s_

        if (present(fmt)) then
            write (s_, "("//fmt//")") r
        else
            write (s_, *) r
        end if
        s = trim(s_)
    end function to_string_real

    !> 将 Fortran 字符串转换为 C 字符串
    pure function f_c_string(f_string) result(c_string)
        character(*), intent(in) :: f_string       !! Fortran 字符串
        character(len(f_string) + 1) :: c_string

        c_string = f_string//c_null_char
    end function f_c_string

    !> 生成一个进度条字符串
    pure function progress_bar(msg, p) result(bar)
        character(*), intent(in) :: msg     !! 进度条描述
        real, intent(in) :: p               !! 当前进度百分比
        character(:), allocatable :: bar    !! 进度条
        integer :: length_, l
        parameter(l=35)

        length_ = len(msg) + l + 11
        allocate (character(length_) :: bar)
        write (bar, "(A,A2,A35,A2,F5.1,A2)") msg, &
            " [", repeat("#", nint(p*l))//repeat(".", l - nint(p*l)), "] ", 100*p, " %"

    end function progress_bar

    !> 不换行输出字符串
    impure subroutine puts(line)
        character(*), intent(in) :: line   !! 输出字符串
        write (*, '(a)', advance='no') line
    end subroutine puts

    !> 判断字符串是否以某个字符串开头; 如果 back 为真, 则从后往前查找 (ends_with)
    pure logical function starts_with(string, substring, back) result(ok)
        character(*), intent(in) :: string      !! 主字符串
        character(*), intent(in) :: substring   !! 子字符串
        logical, intent(in), optional :: back   !! 是否从后面开始比较
        logical :: back_

        if (present(back)) then
            back_ = back
        else
            back_ = .false.
        end if

        if (back_) then
            ok = index(string, substring, .true.) + len(substring) - 1 == len(string)
        else
            ok = index(string, substring) == 1
        end if

    end function starts_with

    !> Error message
    subroutine error_stop_msg(msg, procedure, no_backtrace)
        character(*), intent(in) :: msg         !! error message
        character(*), intent(in) :: procedure   !! procedure name
        logical, intent(in), optional :: no_backtrace     !! no backtrace flag, default is .false.
        logical :: no_backtrace_
        character(:), allocatable :: error_procedure_msg

        if (present(no_backtrace)) then
            no_backtrace_ = no_backtrace
        else
            no_backtrace_ = .false.
        end if

        write (*, '(2a)') 'ERROR MESSAGE: ', msg
        error_procedure_msg = 'PROCEDURE: '//procedure

        if (no_backtrace_) then
            stop error_procedure_msg
        else
            error stop error_procedure_msg
        end if

    end subroutine error_stop_msg

    !> Error code
    subroutine error_stop_code(code, procedure, no_backtrace)
        integer, intent(in) :: code             !! error code
        character(*), intent(in) :: procedure   !! procedure name
        logical, intent(in), optional :: no_backtrace     !! no backtrace flag, default is .false.
        logical :: no_backtrace_
        character(:), allocatable :: error_procedure_msg

        if (present(no_backtrace)) then
            no_backtrace_ = no_backtrace
        else
            no_backtrace_ = .false.
        end if

        write (*, '(a,i0)') 'ERROR CODE: ', code
        error_procedure_msg = 'PROCEDURE: '//procedure

        if (no_backtrace_) then
            stop error_procedure_msg
        else
            error stop error_procedure_msg
        end if

    end subroutine error_stop_code

    !> Inquire whether a file/directory exists.
    impure elemental logical function is_exist(file_name, dir)
        character(*), intent(in) :: file_name !! file name
        logical, intent(in), optional :: dir  !! directory
#if defined __INTEL_COMPILER
        if (present(dir)) then
            if (dir) then
                inquire (directory=file_name, exist=is_exist)
                return
            end if
        end if
#endif
        inquire (file=file_name, exist=is_exist)

    end function is_exist

    !> Make directory
    impure elemental subroutine mkdir(file_name)
        character(*), intent(in) :: file_name  !! file name
        integer :: exitstat
        character(40) :: cmdmsg
        if (is_windows()) then
            call execute_command_line("md "//windows_path(file_name), exitstat=exitstat)
        else
            call execute_command_line("mkdir -p "//file_name, exitstat=exitstat)
        end if
    end subroutine mkdir

    !> Remove directory
    impure elemental subroutine rmdir(file_name)
        character(*), intent(in) :: file_name   !! file name
        integer :: exitstat
        if (is_windows()) then
            call execute_command_line("rd /s /q "//windows_path(file_name), exitstat=exitstat)
        else
            call execute_command_line("rm -rf "//file_name, exitstat=exitstat)
        end if
    end subroutine rmdir

    !> Remove file
    impure elemental subroutine remove(file_name)
        character(*), intent(in) :: file_name !! file name
        integer tmp_unit
        if (is_exist(file_name)) then
            open (newunit=tmp_unit, file=file_name, status='old')
            close (tmp_unit, status='delete')
        end if
    end subroutine remove

    !> Judge if the system is windows
    logical function is_windows()
        character(16) :: os_name
        logical, save :: is_windows_ = .false.
        logical, save :: is_first_run = .true.

        if (is_first_run) then
            call getenv("OS", os_name) ! GNU extension, IFORT extension
            is_windows_ = trim(os_name) == "Windows_NT"
            is_first_run = .false.
            is_windows = is_windows_
        else
            is_windows = is_windows_
        end if

    end function is_windows

    !> Get current working directory
    function pwd()
        character(:), allocatable :: pwd
        character(256) :: pwd_
        call getcwd(pwd_) ! GNU extension, IFORT extension
        pwd = trim(pwd_)
    end function pwd

    ! gfortran has an rename intrinsic function, but here gived the another.
    !> Move file or rename file
    subroutine rename(old_file, new_file)
        character(*), intent(in) :: old_file !! old file name
        character(*), intent(in) :: new_file !! new file name
        integer :: exitstat
        if (is_windows()) then
            call execute_command_line('ren "'//windows_path(old_file)//'" "'// &
                                      new_file(scan(new_file, "/\", .true.) + 1:)//'"', &
                                      exitstat=exitstat)
        else
            call execute_command_line("mv "//old_file//" "//new_file, exitstat=exitstat)
        end if
    end subroutine rename

    !> Counts the line number of file
    integer function countlines(file_name)
        character(*), intent(in) :: file_name !! file name
        integer :: istat, iunit
        open (newunit=iunit, file=file_name, status='old')
        countlines = 0
        do
            read (iunit, *, iostat=istat)
            if (is_iostat_end(istat)) exit
            countlines = countlines + 1
        end do
        close (iunit)
    end function countlines

    !> Replace file system separator for Windows
    function windows_path(path) result(winpath)
        character(*), intent(in) :: path
        character(len(path)) :: winpath
        integer :: ik

        do concurrent(ik=1:len(path))
            if (path(ik:ik) == "/") then
                winpath(ik:ik) = "\"
            else
                winpath(ik:ik) = path(ik:ik)
            end if
        end do

    end function windows_path

    !> 启动计时器
    impure subroutine tic()
        call system_clock(time_save)
    end subroutine tic

    !> 停止计时器并返回时间（秒计时）
    impure subroutine toc(t)
        real, intent(out), optional :: t !! 时间（秒计时）
        integer :: time_now, time_rate

        call system_clock(time_now, time_rate)

        associate (dt => real(time_now - time_save)/time_rate)

            if (present(t)) then
                t = dt
            else
                write (*, "(a,g0.3,a)") 'Time elapsed: ', dt, " s"
            end if

        end associate

    end subroutine toc

    !> 获得当前日期或时间
    impure function time_stamp() result(t)
        character(:), allocatable :: t
        character(len=8) :: datstr
        character(len=10) :: timstr

        call date_and_time(datstr, timstr)

        t = datstr(1:4)//"-"// &
            datstr(5:6)//"-"// &
            datstr(7:8)//" "// &
            timstr(1:2)//":"// &
            timstr(3:4)//":"// &
            timstr(5:10)

    end function time_stamp

    !> 判断是否为闰年
    logical pure function is_leap(year)
        integer, intent(in) :: year !! 检查的年份

        is_leap = (mod(year, 4) == 0 .and. mod(year, 100) /= 0) .or. mod(year, 400) == 0

    end function is_leap

    !> constructs the identity matrix.
    pure function eye(m, n) result(a)
        integer, intent(in) :: m            !! number of rows
        integer, intent(in), optional :: n  !! number of columns
        integer, allocatable :: a(:, :)
        integer :: i, j, n_

        if (present(n)) then
            n_ = n
        else
            n_ = m
        end if

        allocate (a(m, n_))

        do concurrent(j=1:n_, i=1:m)
            if (i == j) then
                a(i, j) = 1
            else
                a(i, j) = 0
            end if
        end do
    end function eye

    !> constructs a vector of integers from start to end with step.
    pure function arange_integer(start, end, step) result(v)
        integer, intent(in) :: start
        integer, intent(in), optional :: end, step
        integer, allocatable :: v(:)

        integer start_, end_, step_
        integer i

        if (present(end)) then
            start_ = start
            end_ = end
        else
            start_ = 1
            end_ = start
        end if

        if (present(step)) then
            step_ = step
        else
            step_ = 1
        end if

        if (step_ /= 0) then
            step_ = sign(step_, end_ - start_)
        else
            step_ = sign(1, end_ - start_)
        end if

        allocate (v((end_ - start_)/step_ + 1))
        v = [(i, i=start_, end_, step_)]
    end function arange_integer

    !> constructs a vector of reals from start to end with step.
    pure function arange_real(start, end, step) result(v)
        real, intent(in) :: start
        real, intent(in), optional :: end, step
        real, allocatable :: v(:)

        real start_, end_, step_
        integer i

        if (present(end)) then
            start_ = start
            end_ = end
        else
            start_ = 1.0
            end_ = start
        end if

        if (present(step)) then
            step_ = step
        else
            step_ = 1.0
        end if

        if (step_ /= 0.0) then
            step_ = sign(step_, end_ - start_)
        else
            step_ = sign(1.0, end_ - start_)
        end if

        allocate (v(floor((end_ - start_)/step_) + 1))
        v = [(start_ + (i - 1)*step_, i=1, size(v), 1)]
    end function arange_real

    !> constructs a vector of the diagonal elements of a matrix.
    pure function diag_real_rank2(a) result(v)
        real, intent(in) :: a(:, :)
        real :: v(min(size(a, 1), size(a, 2)))
        integer i

        do concurrent(i=1:size(v))
            v(i) = a(i, i)
        end do
    end function diag_real_rank2

    !> constructs a matrix with the diagonal elements of a vector.
    pure function diag_real_rank1(v) result(a)
        real, intent(in) :: v(:)
        real :: a(size(v), size(v))
        integer :: i, j

        do concurrent(j=1:size(v), i=1:size(v))
            if (i == j) then
                a(i, j) = v(i)
            else
                a(i, j) = 0.0
            end if
        end do
    end function diag_real_rank1

    !> constructs a vector of n linearly spaced numbers from start to end.
    pure function linspace_real(start, end, n) result(v)
        real, intent(in) :: start, end
        integer, intent(in) :: n
        real :: v(max(n, 0))

        real step
        integer i

        if (n <= 0) then
            return
        elseif (n == 1) then
            v(1) = start
            return
        end if

        step = (end - start)/(n - 1)
        do concurrent(i=1:n)
            v(i) = start + (i - 1)*step
        end do

    end function linspace_real

    !> returns a vector of n linearly spaced integers.
    pure function linspace_integer(start, end, n) result(v)
        integer, intent(in) :: start, end
        integer, intent(in) :: n
        integer v(max(n, 0))

        integer step
        integer i

        if (n <= 0) then
            return
        elseif (n == 1) then
            v(1) = start
            return
        end if

        step = (end - start)/(n - 1)
        do concurrent(i=1:n)
            v(i) = start + (i - 1)*step
        end do
    end function linspace_integer

    !> returns the lower triangular part of a.
    pure function tril(a, k) result(b)
        real, intent(in) :: a(:, :)
        integer, intent(in), optional :: k
        real :: b(size(a, 1), size(a, 2))
        integer :: k_, i, j

        if (present(k)) then
            k_ = k
        else
            k_ = 0
        end if

        do concurrent(j=1:size(a, 2), i=1:size(a, 1))
            if (i - k_ < j) then
                b(i, j) = 0.0
            else
                b(i, j) = a(i, j)
            end if
        end do
    end function tril

    !> returns the upper triangular part of a.
    pure function triu(a, k) result(b)
        real, intent(in) :: a(:, :)
        integer, intent(in), optional :: k
        real :: b(size(a, 1), size(a, 2))
        integer :: k_, i, j

        if (present(k)) then
            k_ = k
        else
            k_ = 0
        end if

        do concurrent(j=1:size(a, 2), i=1:size(a, 1))
            if (i - k_ < j) then
                b(i, j) = a(i, j)
            else
                b(i, j) = 0.0
            end if
        end do
    end function triu

    !> returns a tridiagonal matrix.
    pure function tri(m, n, k) result(a)
        integer, intent(in) :: m
        integer, intent(in), optional :: n, k
        integer, allocatable :: a(:, :)
        integer :: n_, k_, i, j

        if (present(n)) then
            n_ = n
        else
            n_ = m
        end if

        if (present(k)) then
            k_ = k
        else
            k_ = 0
        end if

        allocate (a(m, n_))
        do concurrent(j=1:n_, i=1:m)
            if (i - k_ < j) then
                a(i, j) = 0
            else
                a(i, j) = 1
            end if
        end do

    end function tri

end module forlab_m
